#!/usr/bin/env python3
'''
正在修复bug

验证失败，虽然训练平均分数上升，但是貌似只能到将近30左右（并且测试分数也是很奇怪的稳定在31分就没有继续上升了，也没有下降，一直都是31分），就不在上升
不知道原因是什么，有待排查？可能也是因为训练的时间不够就
可以跑一下训练好的31分最好模型看看是什么情况
不用验证，因为其他模型有跑到300多分的
'''
import gymnasium as gym
import ptan
import argparse
from tensorboardX import SummaryWriter

import torch
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim
import torch.multiprocessing as mp
import os

from lib import common

GAMMA = 0.99
LEARNING_RATE = 5e-4
ENTROPY_BETA = 0.01

REWARD_STEPS = 4
CLIP_GRAD = 0.1

PROCESSES_COUNT = 4
NUM_ENVS = 15

GRAD_BATCH = 64 # 控制用来训练的梯度数据缓存中长度，这里之所以将GRAD_BATCH设置的比较大，是因为损失计算和反向传播是非常繁重的任务，这样设置可以提高效率；如果设置的比较小，梯度传播很频繁，会导致因为同步梯度导致网络权重被锁，效率低下
ENV_NAME = "ALE/Amidar-v5"
NAME = "Amidar"
REWARD_BOUND = 1000
TRAIN_BATCH = 4

class RewardWrapper(gym.Wrapper):
    def __init__(self, env):
        super(RewardWrapper, self).__init__(env)
        self.previous_lives = None  # 用于跟踪智能体的生命值

    def reset(self, **kwargs):
        # 重置环境并初始化生命值
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info['lives']
        return obs, info

    def step(self, action):
        obs, reward, done, trunc, info = self.env.step(action)

        # 自定义奖励机制
        custom_reward = self.custom_reward(info, reward, done, action)

        return obs, custom_reward, done, trunc, info

    def custom_reward(self, info, original_reward, done, action):
        custom_reward = 0

        # 基于原始游戏的奖励，如击杀外星人时的正向奖励
        custom_reward += original_reward

        # 获取当前生命值
        current_lives = info['lives']

        # 检查是否死亡并给予惩罚
        if current_lives < self.previous_lives:
            custom_reward -= 20  # 死亡惩罚
        self.previous_lives = current_lives

        # 存活奖励（每一步给予小奖励）
        if not done:
            custom_reward += 0.1  # 每步存活的奖励

        # 如果有其他道具收集机制，你可以在这里添加道具收集奖励
        # if self.item_collected():
        #     custom_reward += 5  # 收集道具奖励

        return custom_reward

'''
    在 Amidar 游戏中，奖励系统基于玩家在迷宫中着色线条和捕捉敌人（在某些条件下被转化为小鸡）。以下是详细的奖励机制：

奖励机制：
着色线条：

着色迷宫内的每一条水平线，根据线条的长度，得分从 1 到 6 分不等。最短的线条为 1 分，最长的线条为 6 分。
着色每一条垂直线，得 1 分。
着色完整的方格：

着色完成一个方格的四条边（即完整着色一个方格），得 50 分。
抓住小鸡：

在特定情况下，当完成着色四个角落的方格时，敌人会变成小鸡。抓住小鸡可以得到 100 分 的奖励。
游戏难度及回合：
游戏从第1关（5个敌人，速度最慢）开始，并在不同关卡交替使用不同的角色（如“战士和大猩猩”或“猪和油漆滚子”）。
游戏的速度随着回合的增加而加快，最高可达到第6关，该关卡有6个敌人并且速度最快。
奖励总结：
基础奖励：着色迷宫内的线条和方格可持续获得少量积分，完成方格着色获得较大的积分。
特殊奖励：捕捉小鸡可以获得100分的额外奖励。
'''

class BreakoutPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(BreakoutPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        
        # 每帧惩罚
        # reward += self.frame_penalty
        
        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward += self.life_loss_penalty
            self.previous_lives = current_lives
        
        return obs, reward, done, truncated, info
    

def wrap_dqn(env, stack_frames=4, episodic_life=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)
    # 跳帧包装器
    # env = ptan.common.wrappers.MaxAndSkipEnv(env, skip=4)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = ptan.common.wrappers.FireResetEnv(env)
    env = common.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = BreakoutPenaltyWrapper(env)
    return env

def make_env():
    # # 创建环境，将不同环境统一成统一操作
    # 比如有些游戏只有一条命，有些游戏有多条命
    return wrap_dqn(gym.make(ENV_NAME, obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)


def grads_func(proc_name, net, device, train_queue):
    '''
    梯度传播子进程执行函数
    '''
    # 每个子进程创建NUM_ENVS个环境
    envs = [make_env() for _ in range(NUM_ENVS)]
    # 创建策略代理
    agent = ptan.agent.PolicyAgent(lambda x: net(x)[0], device=device, apply_softmax=True)
    # 创建经验源，经验源本身支持多环境
    exp_source = ptan.experience.ExperienceSourceFirstLast(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)

    batch = []
    frame_idx = 0
    # 为记录统计设置名字，设置为每个子进程的名字
    writer = SummaryWriter(comment=proc_name)

    # todo 完善RewardTracker注释 作用统计近100次游戏的奖励，如果奖励达标则退出游戏
    with common.RewardTracker(writer, stop_reward=REWARD_BOUND) as tracker:
        # todo 完善TBMeanTracker注释，作用将最近的10步游戏记录到tensorboard中
        with ptan.common.utils.TBMeanTracker(writer, batch_size=100) as tb_tracker:
            # 从经验源中获取经验样本
            for exp in exp_source:
                frame_idx += 1
                # 获取经验源中所有的回报
                # pop_total_rewards存在返回值，那么表示一次游戏
                # 已经结束
                new_rewards = exp_source.pop_total_rewards()
                # 判断游戏目前运行获取的奖励是否已经达到标准，达到了就退出训练
                if new_rewards and tracker.reward(new_rewards[0], frame_idx):
                    break

                # 判断目前已经收集到的经验训练数据数是否已经达到目标
                batch.append(exp)
                if len(batch) < GRAD_BATCH:
                    continue

                # 解压出训练数据
                states_v, actions_t, vals_ref_v = \
                    common.unpack_batch_a3c(batch, net, last_val_gamma=GAMMA**REWARD_STEPS, device=device)
                batch.clear()

                net.zero_grad()
                # 得到动作预测，当前状态下的Q值
                logits_v, value_v = net(states_v)
                # 训练网络预测Q值，预测出来的值和bellman计算出来的Q值应该接近
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                # 计算每个动作的概率
                log_prob_v = F.log_softmax(logits_v, dim=1)
                adv_v = vals_ref_v - value_v.detach()
                log_prob_actions_v = adv_v * log_prob_v[range(GRAD_BATCH), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                prob_v = F.softmax(logits_v, dim=1)
                entropy_loss_v = ENTROPY_BETA * (prob_v * log_prob_v).sum(dim=1).mean()

                loss_v = entropy_loss_v + loss_value_v + loss_policy_v
                loss_v.backward()

                # 将相关的鹅值记录起来
                tb_tracker.track("advantage", adv_v, frame_idx)
                tb_tracker.track("values", value_v, frame_idx)
                tb_tracker.track("batch_rewards", vals_ref_v, frame_idx)
                tb_tracker.track("loss_entropy", entropy_loss_v, frame_idx)
                tb_tracker.track("loss_policy", loss_policy_v, frame_idx)
                tb_tracker.track("loss_value", loss_value_v, frame_idx)
                tb_tracker.track("loss_total", loss_v, frame_idx)

                # gather gradients
                # 梯度裁剪，防止梯度过大或者过小
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                # 将网络参数中每个参数的梯度提取出来
                grads = [param.grad.data.cpu().numpy() if param.grad is not None else None
                         for param in net.parameters()]
                # 放到train_queue共享队列中
                train_queue.put(grads)

    # 这里之所以将None置入队列，是因为主进程判断到队列中有None，说明已经检测到满足了
    # 训练的目标，可以达到目标的奖励值了
    train_queue.put(None)

import numpy as np

def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        obs, _ = env.reset()
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v, _ = net(obs_v)
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes

if __name__ == "__main__":
    # 该方法主要有两个作用:
    #
    # 设置多进程的启动方式(spawn或fork)。默认是spawn。
    # 避免多进程模块的多次导入导致的错误。
    # 关于启动方式:
    #
    # spawn: 子进程将只继承父进程的信号处理机制,而不会继承父进程的句柄、锁等资源。所以spawn方式更安全,也是推荐的默认方式。
    # fork: 子进程通过fork启动,会通过复制父进程获得父进程的全部资源,包括锁、信号处理等,所以可能会导致在多进程中出现问题。
    mp.set_start_method('spawn')
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default="a3c-grad", required=False, help="Name of the run")
    args = parser.parse_args()
    device = "cuda" if args.cuda else "cpu"

    env = make_env()
    test_env = make_env()
    net = common.AtariA2C(env.observation_space.shape, env.action_space.n).to(device)
    # 生命网络net为共享内存，share_memory必须在调用fork前调用
    # 这样net就可以很方便的在每个进程间使用了
    # 使用share_memory会共享网络权重，而不会共享计算图、梯度
    # 在PyTorch中，`share_memory`方法是针对多进程场景设计的，特别是当使用`torch.multiprocessing`时。当你在主进程中创建一个`nn.Module`模型并想在多个子进程中共享这个模型时，使用`share_memory`方法可以将模型的权重存放在共享内存中，从而允许多个子进程访问。
    #
    # 然而，这并不意味着梯度是共享的。当每个子进程对模型进行前向和反向传播时，它们各自会有自己的计算图和梯度存储。这意味着，如果多个进程同时对同一个模型进行反向传播，它们不会在共享内存中累积梯度。
    #
    # 如果你想在多个进程中累积梯度或同步模型参数，你需要使用更高级的同步策略。这种策略在分布式训练场景中更为常见，例如使用`torch.nn.parallel.DistributedDataParallel`。
    #
    # 简而言之，虽然`share_memory`方法允许在多进程中共享模型权重，但梯度并不共享。每个进程都会有自己的计算图和梯度存储。
    net.share_memory()

    step_idx = 0
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10000, gamma=0.9)

    # 创建保存训练过程中数据的路径
    save_path = os.path.join("saves", "a3c-" + args.name)
    os.makedirs(save_path, exist_ok=True)

    # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        step_idx = checkpoint['step_idx']
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        print(f"Current learning rate: {scheduler.get_last_lr()[0]}")

        print("加载模型成功")

    # 创建子进程共享队列，最大存储的数量为子进程的数量
    train_queue = mp.Queue(maxsize=PROCESSES_COUNT)
    data_proc_list = []
    for proc_idx in range(PROCESSES_COUNT):
        # 创建子进程
        # 为啥梯度同步需要创建进程名字
        # 因为每个梯度同步就是每个子进程内部都包含一个自己训练网路网络的流程
        # 由于要记录summary，所以需要区分每个进程的记录名字
        proc_name = "-a3c-grad_" + NAME + "_" + args.name + "#%d" % proc_idx
        data_proc = mp.Process(target=grads_func, args=(proc_name, net, device, train_queue))
        data_proc.start()
        data_proc_list.append(data_proc)

    batch = []
    # 梯度缓冲区，接受子进程中返回的梯度信息
    grad_buffer = None
    best_reward = -float('inf')

    try:
        while True:
            # 从梯度进程共享缓冲区提取梯度信息
            # get时会等待，知道有数据，每次get只会获取一个梯度信息
            # 而根据返回值，每次返回的梯度信息均包含每个网路节点的梯度信息
            # 顺序则按照parameters返回值而定
            train_entry = train_queue.get()
            if train_entry is None:
                break

            step_idx += 1

            # 接收子进程中的梯度信息
            if grad_buffer is None:
                # 如果grad_buffer是空，则直接赋值
                grad_buffer = train_entry
            else:
                # 如果grad_buffer不为空，则需要将gard_buffer和train_entry中的梯度相加汇总（符合书上）
                # 而多个子进程的梯度信息也只是简单的加起来
                for tgt_grad, grad in zip(grad_buffer, train_entry):
                    tgt_grad += grad

            # 判断是否需要同步梯度到网路权重
            if step_idx % TRAIN_BATCH == 0:
                for param, grad in zip(net.parameters(), grad_buffer):
                    # 这里是将子进程汇总的梯度设置到对应的权重的梯度上
                    param.grad = torch.FloatTensor(grad).to(device)

                # 梯度裁剪
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                # 梯度更新，更新网路权重
                optimizer.step()
                scheduler.step()
                grad_buffer = None

            if step_idx % 200 == 0:
                # Test the model
                test_reward = test_model(test_env, net, device=device, episodes=2)
                print(f"Test reward: {test_reward:.2f}")

                # Save the model if it is the best so far
                if test_reward > best_reward:
                    best_reward = test_reward
                    model_path = os.path.join(save_path, f"best_model_{best_reward:.2f}.pth")
                    torch.save(net.state_dict(), model_path)
                    print(f"New best model saved with reward {best_reward:.2f}")

            if step_idx % 100 == 0:
                checkpoint = {
                    "net": net.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "step_idx": step_idx,
                    'scheduler_state_dict': scheduler.state_dict()
                }
                common.save_checkpoints(step_idx, checkpoint, save_path, "a3c", keep_last=5)
    finally:
        # 之所以这里要额外进行这个动作，是为了防止僵尸进程占据GPU资源
        # 因为这里是采用多进程的方式，主进程退出子进程很可能不会退出
        for p in data_proc_list:
            p.terminate()
            p.join()
